import re
import csv
import os
import dgl
import torch
import multiprocessing

from collections import defaultdict
from parse_lib import parse_lib_pin_info 

def parse_sdc(sdc_file_path):
    pass

def clean_name(name):
    """去除 DEF 文件中的转义符"""
    return name.replace('\/', '/').replace('\[', '[').replace('\]', ']')

def calculate_pin_distances(unit_x, unit_y, unit_width, unit_height, orientation, pin_x, pin_y, die_width, die_height):
    """
    计算单元引脚相对于die四个边界的距离

    参数：
    - unit_x, unit_y: 单元左下角在die中的坐标
    - unit_width, unit_height: 单元的宽度和高度
    - orientation: 单元的摆放方向
    - pin_x, pin_y: 单元内引脚相对于单元左下角的坐标
    - die_width, die_height: die的宽度和高度
    """
    if orientation == "N":  # North，不旋转
        pin_global_x = unit_x + pin_x
        pin_global_y = unit_y + pin_y
    elif orientation == "S":  # South，旋转180度
        pin_global_x = unit_x + (unit_width - pin_x)
        pin_global_y = unit_y + (unit_height - pin_y)
    elif orientation == "E":  # East，旋转90度顺时针
        pin_global_x = unit_x + pin_y
        pin_global_y = unit_y + (unit_width - pin_x)
    elif orientation == "W":  # West，旋转90度逆时针
        pin_global_x = unit_x + (unit_height - pin_y)
        pin_global_y = unit_y + pin_x
    elif orientation == "FN":  # Flip North，左右翻转
        pin_global_x = unit_x + (unit_width - pin_x)
        pin_global_y = unit_y + pin_y
    elif orientation == "FS":  # Flip South，左右翻转
        pin_global_x = unit_x + pin_x
        pin_global_y = unit_y + (unit_height - pin_y)
    elif orientation == "FE":  # Flip East，翻转后顺时针旋转90度
        pin_global_x = unit_x + pin_y
        pin_global_y = unit_y + pin_x
    elif orientation == "FW":  # Flip West，翻转后逆时针旋转90度
        pin_global_x = unit_x + (unit_height - pin_y)
        pin_global_y = unit_y + (unit_width - pin_x)
    else:
        raise ValueError("Invalid orientation. It must be one of ['N', 'S', 'E', 'W', 'FN', 'FS', 'FE', 'FW'].")
    
    # 计算引脚到die四个边界的距离
    left_distance = pin_global_x
    right_distance = die_width - pin_global_x
    bottom_distance = pin_global_y
    top_distance = die_height - pin_global_y

    return left_distance, right_distance, top_distance, bottom_distance

def parse_timing_rpt(timing_reports_dir):
    """
    解析指定目录下的所有 Timing 报告文件，提取 Endpoint 信息, startPoint信息
    """
    endpoint_pins = set()
    beginpoint_pins = set()
    
    for filename in os.listdir(timing_reports_dir):
        if "timing" in filename:
            rpt_path = os.path.join(timing_reports_dir, filename)
            with open(rpt_path, 'r') as timing_file:
                for line in timing_file:
                    endpoint_match = re.search(r'Endpoint:\s+([\w\d_/\[\]]+)', line)
                    # 查找 Endpoint
                    if endpoint_match:
                        endpoint_pins.add(endpoint_match.group(1))
                    
                    # 查找 Beginpoint
                    beginpoint_match = re.search(r'Beginpoint:\s+([\w\d_/\[\]]+)', line)
                    if beginpoint_match:
                        beginpoint_pins.add(beginpoint_match.group(1))
    
    return endpoint_pins, beginpoint_pins

#特殊情况！
#   (CELL
#     (CELLTYPE  "XNOR2xp5_ASAP7_75t_L")
#     (INSTANCE  U2379)
#       (DELAY
# 	(ABSOLUTE
# 	(COND ~B (IOPATH A Y  (9.000::9.700) (6.900::7.400)))
# 	(COND B (IOPATH A Y  (10.700::11.100) (14.900::15.700)))
# 	(COND ~A (IOPATH B Y  (11.800::12.900) (6.600::7.000)))
# 	(COND A (IOPATH B Y  (12.300::12.800) (15.400::16.300)))
# 	)
#       )
#   )

#   (CELL
#     (CELLTYPE  "DFFHQNx2_ASAP7_75t_L")
#     (INSTANCE  sa13_reg\[6\])
#       (DELAY
# 	(ABSOLUTE
# 	(IOPATH (posedge CLK) QN  (59.700::61.000) (59.100::60.300))
# 	)
#       )
#       (TIMINGCHECK
# 	(SETUPHOLD (posedge D) (posedge CLK) (6.458::4.750) (-1.286::-1.152))
# 	(SETUPHOLD (negedge D) (posedge CLK) (-2.861::1.887) (13.724::12.374))
# 	(WIDTH (COND D (posedge CLK)) (41.492::44.614))
# 	(WIDTH (COND D (negedge CLK)) (39.111::41.895))
# 	(WIDTH (COND ~D (posedge CLK)) (41.492::44.614))
# 	(WIDTH (COND ~D (negedge CLK)) (39.111::41.895))
#       )
#   )

def parse_sdf_net_section(sdf_file_path):
    net_connections = set()
    sdf_info = {}
    with open(sdf_file_path, 'r') as file:
        sdf_content = file.read() 
        # 正则表达式匹配 INTERCONNECT 行
        net_delay_pattern = r"\(INTERCONNECT\s+([^\s]+)\s+([^\s]+)\s+\(([\d\.\:]+)\)\s+\(([\d\.\:]+)\)\)"

        for line in sdf_content.splitlines():
            net_delay_match = re.match(net_delay_pattern, line.strip())
            if net_delay_match:
                source_pin = clean_name(net_delay_match.group(1))
                target_pin = clean_name(net_delay_match.group(2))  
                net_delay_1 = net_delay_match.group(3) 
                net_delay_2 = net_delay_match.group(4) 

                sdf_info[source_pin] = {
                    'direction': 'output',
                    'FF_rise_net_delay_to_root_pin': 0.000,
                    'SS_rise_net_delay_to_root_pin': 0.000,
                    'FF_fall_net_delay_to_root_pin': 0.000,
                    'SS_fall_net_delay_to_root_pin': 0.000
                }
                sdf_info[target_pin] = {
                        'direction': 'input',
                        'FF_rise_net_delay_to_root_pin': float(net_delay_1.split("::")[0]),  
                        'SS_rise_net_delay_to_root_pin': float(net_delay_1.split("::")[1]), 
                        'FF_fall_net_delay_to_root_pin': float(net_delay_2.split("::")[0]), 
                        'SS_fall_net_delay_to_root_pin': float(net_delay_2.split("::")[1]) 
                }
                net_connections.add((source_pin, target_pin))

    return sdf_info, net_connections

def parse_sdf_cell_connections(sdf_file_path):
    cell_connections = set()
    instance_name = None

    with open(sdf_file_path, "r") as f:
        sdf_content = f.read()
    # 修改正则表达式以匹配包含括号的引脚名称
    iopath_pattern = re.compile(r"\(IOPATH\s+([^\s\(\)]+|\([^\)]+\))\s+([^\s\(\)]+|\([^\)]+\))\s+\(([^)]+)::([^)]+)\)\s+\(([^)]+)::([^)]+)\)\)")
    
    for line in sdf_content.splitlines():
        line = line.strip()
        
        # 判断是否是 CELL 开头
        if line.startswith("(CELL"):
             instance_name = None  # 重置当前实例名
        
        # 匹配 INSTANCE 行
        elif line.startswith("(INSTANCE"):
            #排除(INSTANCE)
            parts = line.split()
            if len(parts) > 1:
                instance_name = clean_name(parts[1].strip(")"))  # 获取 INSTANCE 名称
            else:
                instance_name = None
        
        # 匹配 IOPATH 行
        elif "(IOPATH" in line and instance_name:
            match = iopath_pattern.search(line)
            if match:
                src_pin, dst_pin = match.group(1), match.group(2)
                
                clk_posedge_pattern = r'\(?(posedge|negedge)\s+(\w+)\)?'
                # 判断是否是时钟信号形式（pos/neg edge）
                if "posedge" in src_pin or "negedge" in src_pin:
                    match = re.search(clk_posedge_pattern, src_pin)
                    if match:
                        src_pin = match.group(2)  # 时钟源引脚
                    else:
                        print("ERROR : cell sdf")

                if "posedge" in dst_pin or "negedge" in dst_pin:
                    
                    match = re.search(clk_posedge_pattern, dst_pin)
                    if match:
                        dst_pin = match.group(2)  # 时钟目标引脚
                    else:
                        print("ERROR : cell sdf")

                cell_connections.add((f"{instance_name}/{src_pin}", f"{instance_name}/{dst_pin}"))
    
    return cell_connections

def parse_lef_file(lef_file_path):
    """
    解析一个 LEF 文件，提取所有的 cell 和对应的 pin 信息。
    
    参数:
    lef_file_path: LEF 文件的路径
    
    返回:
    一个字典，键是 cell 名称，值是另一个字典，存储每个 pin 的信息。
    """
    cell_data = {}
    inside_macro = False
    inside_pin = False
    current_pin_name = None
    current_cell_name = None
    
    with open(lef_file_path, 'r') as file:
        lines = file.readlines()

    for line in lines:
        line = line.strip()
        
        # 进入一个新的宏定义
        if line.startswith("MACRO "):
            inside_macro = True
            current_cell_name = line.split()[1].strip()
            cell_data[current_cell_name] = {'pins': {}}
            continue  # 继续读取下一行
        
        if inside_macro:
            # 找到SIZE行并提取cell的尺寸
            if line.startswith("SIZE"):
                size_match = re.match(r"SIZE\s+([0-9\.]+)\s+BY\s+([0-9\.]+)\s*;", line)
                if size_match:
                    cell_width = float(size_match.group(1))
                    cell_height = float(size_match.group(2))
                    cell_data[current_cell_name]['width'] = cell_width
                    cell_data[current_cell_name]['height'] = cell_height

            # 解析PIN信息
            if line.startswith("PIN "):
                inside_pin = True
                current_pin_name = line.split()[1].strip()
                cell_data[current_cell_name]['pins'][current_pin_name] = {'rect': None}
            
            if inside_pin and current_pin_name:
                # 找到PORT行，跳过
                if line.startswith("PORT"):
                    continue  # 跳过PORT行
                
                # 找到RECT并提取坐标
                if line.startswith("RECT"):
                    rect_match = re.match(r"RECT\s*(-?\d+\.?\d*)\s*(-?\d+\.?\d*)\s*(-?\d+\.?\d*)\s*(-?\d+\.?\d*)\s*;", line)
                    if rect_match:
                        pin_x1 = float(rect_match.group(1))
                        pin_y1 = float(rect_match.group(2))
                        pin_x2 = float(rect_match.group(3))
                        pin_y2 = float(rect_match.group(4))
                        cell_data[current_cell_name]['pins'][current_pin_name]['rect'] = (pin_x1, pin_y1, pin_x2, pin_y2)

            # 结束PIN定义
            if inside_pin and line.startswith("END " + current_pin_name):
                inside_pin = False
                current_pin_name = None
            
            # 结束MACRO定义
            if line.startswith("END MACRO"):
                inside_macro = False  # 结束当前宏的解析

    return cell_data

def parse_lef_directory(lef_directory):
    """
    解析目录下的所有 LEF 文件，返回所有 cell 和 pin 的信息。
    
    参数:
    lef_directory: LEF 文件所在的目录
    
    返回:
    一个字典，包含所有 cell 的宽度、高度和 pin 的矩形范围。
    """
    all_cells_data = {}

    for lef_file in os.listdir(lef_directory):
        if lef_file.endswith(".lef"):
            lef_file_path = os.path.join(lef_directory, lef_file)
            cell_data = parse_lef_file(lef_file_path)
            all_cells_data.update(cell_data)  # 将当前文件的数据合并到总数据中
    
    return all_cells_data

def parse_def_die_size(def_file_path):
    """
    给定def content以输出die_width和die_height
    """
    with open(def_file_path, 'r') as file:
        def_content = file.read()

    diearea_pattern = r'DIEAREA\s+\(\s*([0-9]+)\s+([0-9]+)\s*\)\s+\(\s*([0-9]+)\s+([0-9]+)\s*\)\s*;'
    diearea_match = re.search(diearea_pattern, def_content)
    if diearea_match:
        die_x1, die_y1, die_x2, die_y2 = map(int, diearea_match.groups())
        die_width = die_x2 - die_x1
        die_height = die_y2 - die_y1
    else:
        die_width, die_height = None, None

    return die_width, die_height

def parse_def_component_section(def_file_path):
    cell_info = {}  
    inside_components_section = False
    
    with open(def_file_path, 'r') as file:
        for line in file:
            line = line.strip()

            # 解析 COMPONENTS 部分的开始和结束
            if line.startswith("COMPONENTS"):
                inside_components_section = True
                continue
            if inside_components_section and line.startswith("END COMPONENTS"):
                inside_components_section = False

            # 解析 COMPONENTS 中的每个 cell
            if inside_components_section:
                # 匹配并提取组件信息
                match = re.match(r"^- (\S+) (\S+) .* PLACED \(\s*(\d+)\s+(\d+)\s*\) (\w+)", line)
                if match:
                    cell_name = clean_name(match.group(1))  # 处理 cell_name
                    cell_type = clean_name(match.group(2))  # 处理 cell_type
                    cell_x_pos = int(match.group(3))  # 单元左下角的 x 坐标
                    cell_y_pos = int(match.group(4))  # 单元左下角的 y 坐标
                    orientation = match.group(5)  # 单元的方向

                    # 将信息存入字典，cell_name 为键
                    cell_info[cell_name] = {
                        "cell_type": cell_type,
                        "cell_x_pos": cell_x_pos,
                        "cell_y_pos": cell_y_pos,
                        "orientation": orientation
                    }

    return cell_info

def parse_def_net_section(def_file_path):
    net_info = {}  
    inside_nets_section = False
    
    with open(def_file_path, 'r') as file:
        for line in file:
            line = line.strip()

            if line.startswith("NETS"):
                inside_nets_section = True
                continue
            if inside_nets_section and line.startswith("END NETS"):
                break

            # 解析 Net Name
            net_match = re.match(r'\-\s+([\w\d_]+(?:\[[0-9]+\])?)', line)
            if net_match:
                current_net = net_match.group(1)
                continue

            if inside_nets_section:
                # 提取 (cell_name, pin_name)
                matches = re.findall(r'\(\s*([\w\d_\\\/\[\]]+)\s+([\w\d_\\\/\[\]]+)\s*\)', line)
                
                for cell_name, pin_name in matches:
                    cell_name_clean = clean_name(cell_name)

                    cell_type = None
                    # 特殊情况
                    if cell_name_clean == "PIN":
                        pin_full_name = pin_name
                        cell_type = "PIN"
                    else:
                        pin_full_name = f"{cell_name_clean}/{clean_name(pin_name)}"

                    net_info[pin_full_name] = {
                        "cell_name" : cell_name_clean,
                        "pin_name" : pin_name,
                        "net_name" : current_net,
                        "cell_type": cell_type 
                    }
                    
        return net_info

#注意这个pin不是cell的pin
def parse_def_pins_section(def_file_path):
    pin_info = {}

    with open(def_file_path, 'r') as file:
        content = file.read() 

        pin_pattern = r"-\s*(\S+)\s*\+ NET\s*(\S+)\s*\+ DIRECTION\s*(\S+)\s*\+ USE\s*(\S+)\s*\+ LAYER\s*(\S+)\s*\(\s*(-?\d+)\s*(-?\d+)\s*\)\s*\(\s*(-?\d+)\s*(-?\d+)\s*\)\s*\+ PLACED\s*\(\s*(-?\d+)\s*(-?\d+)\s*\)\s*(\S+)"

        matches = re.findall(pin_pattern, content)

        for match in matches:
            pin_name = match[0] 
            net_name = match[1]  
            direction = match[2]  
            use = match[3]  
            layer = match[4]  
            size_x1 = int(match[5])
            size_y1 = int(match[6])  
            size_x2 = int(match[7])
            size_y2 = int(match[8]) 
            x_pos = int(match[9]) 
            y_pos = int(match[10])
            orientation = match[11]  

            pin_info[pin_name] = {
                "net_name": net_name,
                "direction": direction,
                "use": use,
                "layer": layer,
                "size_x1": size_x1,
                "size_y1": size_y1,
                "size_x2": size_x2,
                "size_y2": size_y2,
                "x_pos": x_pos,
                "y_pos": y_pos,
                "orientation": orientation,
                "cell_type": "PIN" 
            }

    return pin_info

#for rf model feature
def parse_pre_route_sdf(sdf_file_path):

    slew_dict = {}
    slew_pattern = re.compile(r'\(SLEW\s+(\S+)\s+\(([\d.eE+-]+)::([\d.eE+-]+)\)\s+\(([\d.eE+-]+)::([\d.eE+-]+)\)\)')

    with open(sdf_file_path, 'r') as f:
        for line in f:
            line = line.strip()
            match = slew_pattern.match(line)
            if match:
                pin_name = clean_name(match.group(1))
                try:
                    vals = list(map(float, match.group(2, 3, 4, 5)))
                    if any(v > 1e5  for v in vals):
                        vals = [0.0, 0.0, 0.0, 0.0]
                except:
                    vals = [0.0, 0.0, 0.0, 0.0]
                slew_dict[pin_name] = {
                    'FF_rise_slew':vals[0],
                    'SS_rise_slew':vals[1],
                    'FF_fall_slew':vals[2],
                    'SS_fall_slew':vals[3]
                }
                
    return slew_dict


def parse_pins(def_file, lib_dir, timing_reports_dir, lef_dir, sdf_file_path, pre_route_sdf_path):
    """
    解析 DEF 文件，结合 LIB 和 LEF 文件中的信息，生成 CSV，包含引脚的 Net Name 和 Endpoint 标记。
    """ 

    # 解析 Timing 报告，提取 Endpoint
    endpoint_pins, beginpoint_pins = parse_timing_rpt(timing_reports_dir)

    sdf_file_info ,net_connections= parse_sdf_net_section(sdf_file_path)

    #解析lef files
    lef_files_info = parse_lef_directory(lef_dir)
    
    lib_info = parse_lib_pin_info(lib_dir)
    #die size
    die_width, die_height = parse_def_die_size(def_file)

    def_component_info = parse_def_component_section(def_file)

    def_net_info = parse_def_net_section(def_file)
    
    def_pin_info = parse_def_pins_section(def_file)

    #pre_route sdf
    pre_route_sdf_info = parse_pre_route_sdf(pre_route_sdf_path)

    #引脚特征信息
    pin_info = {}

    pins_no_appear_in_sdf = [] #在def中出现但是没在sdf中出现的pin如 SE
    #SE SI SO并未出现在sdf中注意删除这个引脚

    #注意处理cell type为PIN的特殊情况
    for pin_full_name in def_net_info:

        if pin_full_name in pre_route_sdf_info:
                pre_route_slew_ff_rise = pre_route_sdf_info[pin_full_name]["FF_rise_slew"]
                pre_route_slew_ss_rise = pre_route_sdf_info[pin_full_name]["SS_rise_slew"]
                pre_route_slew_ff_fall = pre_route_sdf_info[pin_full_name]["FF_fall_slew"]
                pre_route_slew_ss_fall = pre_route_sdf_info[pin_full_name]["SS_fall_slew"]

        else:
            pre_route_slew_ff_rise, pre_route_slew_ss_rise, pre_route_slew_ff_fall, pre_route_slew_ss_fall = 0

        #排除不在sdf中的pin
        if pin_full_name in sdf_file_info:
            cell_name = def_net_info[pin_full_name]["cell_name"]   
            cell_type = def_net_info[pin_full_name]["cell_type"]    
            pin_name = def_net_info[pin_full_name]["pin_name"]
            net_name = def_net_info[pin_full_name]["net_name"]
            if not cell_type == 'PIN':
                cell_type = def_component_info[cell_name]["cell_type"]

            cap_data = lib_info.get(cell_type, {}).get(pin_name, {})
            cap_ff_rise = cap_data.get('FF', {}).get('rise', 0.0)
            cap_ff_fall = cap_data.get('FF', {}).get('fall', 0.0)
            cap_ss_rise = cap_data.get('SS', {}).get('rise', 0.0)
            cap_ss_fall = cap_data.get('SS', {}).get('fall', 0.0)

            sdf_pin_data = sdf_file_info.get(pin_full_name,{})

            FF_rise_net_delay_to_root_pin = sdf_pin_data.get('FF_rise_net_delay_to_root_pin', 0.000)
            SS_rise_net_delay_to_root_pin = sdf_pin_data.get('SS_rise_net_delay_to_root_pin', 0.000)
            FF_fall_net_delay_to_root_pin = sdf_pin_data.get('FF_fall_net_delay_to_root_pin', 0.000)
            SS_fall_net_delay_to_root_pin = sdf_pin_data.get('SS_fall_net_delay_to_root_pin', 0.000)

            #注意pin的fanin和fanout的判断和cell的pin的fanin和fanout的判断有所不同
            if cell_type == 'PIN':
                def_pin_info_data = def_pin_info.get(pin_full_name, {})
                pin_x_pos = def_pin_info_data.get("x_pos")
                pin_y_pos = def_pin_info_data.get("y_pos")
                pin_dirction = def_pin_info_data.get("direction")

                is_PI_PO_or_not = 1  #1:is PI or PO ,0 : not

                if pin_x_pos and pin_y_pos:
                    left_distance = pin_x_pos
                    right_distance = die_width - pin_x_pos
                    top_distance = die_height - pin_y_pos
                    bottom_distance = pin_y_pos
                else:
                    left_distance = right_distance = top_distance = bottom_distance = 0

                #is_fanin_or_fanout, 涉及到模型的推理故要保证正确性
                if pin_dirction == "OUTPUT":
                    is_fanin_or_fanout = 0
                elif pin_dirction == "INPUT":
                    is_fanin_or_fanout = 1
                else:
                    print(f"def_pin : {pin_full_name} direction error")
                    is_fanin_or_fanout = None
                    

            else:
                unit_x = def_component_info.get(cell_name, {}).get("cell_x_pos", 0)
                unit_y = def_component_info.get(cell_name, {}).get("cell_y_pos", 0)
                unit_orientation = def_component_info.get(cell_name, {}).get("orientation", "N")

                pin_dirction = lib_info.get(cell_type, {}).get(pin_name, {}).get("direction", None)

                is_PI_PO_or_not = 0

                if pin_dirction == "input":
                    is_fanin_or_fanout = 0
                elif pin_dirction == "output":
                    is_fanin_or_fanout = 1
                else:
                    print(f"cell_pin : {pin_full_name} direction error")
                    is_fanin_or_fanout = None

                lef_cell_info = lef_files_info.get(cell_type, {})
                lef_cell_pin_info = lef_files_info.get(cell_type, {}).get('pins', {}).get(pin_name, {})
                            
                lef_cell_width = lef_cell_info.get('width', None)
                lef_cell_height = lef_cell_info.get('height', None)
                
                if lef_cell_pin_info and lef_cell_pin_info.get('rect') and lef_cell_info:
                    pin_x1, pin_y1, pin_x2, pin_y2 = lef_cell_pin_info['rect']
                    left_distance, right_distance, top_distance, bottom_distance = calculate_pin_distances(
                    unit_x, unit_y, lef_cell_width, lef_cell_height, unit_orientation,
                    pin_x1, pin_y1, die_width, die_height
                                )
                else:
                    left_distance = right_distance = top_distance = bottom_distance = 0

            #增加一步lib解析和sdf解析的fanin fanout的对比防止解析出错
            sdf_direction = sdf_pin_data.get("direction")
            if sdf_direction == "input":
                sdf_is_fanin_or_fanout = 0
            elif sdf_direction == "output":
                sdf_is_fanin_or_fanout = 1
            else:
                print(f"error : {pin_full_name} sdf dirction error")

            if sdf_is_fanin_or_fanout != is_fanin_or_fanout:
                print(f"error : {pin_full_name} sdf direction and direction not right")

            pin_info[pin_full_name] = {
                "cell_type" : cell_type,
                "net_name" : net_name,
                "pin_direction" : pin_dirction,
                "is_fanin_or_fanout" : is_fanin_or_fanout,
                "is_PI_PO_or_not" : is_PI_PO_or_not,
                "ff_rise": cap_ff_rise,
                "ff_fall": cap_ff_fall,
                "ss_rise": cap_ss_rise,
                "ss_fall": cap_ss_fall,
                "is_endpoint": 1 if pin_full_name in endpoint_pins else 0,
                "left_distance":  left_distance,
                "right_distance": right_distance,
                "top_distance": top_distance,
                "bottom_distance": bottom_distance,
                "FF_rise_net_delay_to_root_pin" : FF_rise_net_delay_to_root_pin,
                "SS_rise_net_delay_to_root_pin" : SS_rise_net_delay_to_root_pin,
                "FF_fall_net_delay_to_root_pin" : FF_fall_net_delay_to_root_pin,
                "SS_fall_net_delay_to_root_pin" : SS_fall_net_delay_to_root_pin,
                "FF_rise_slew_pre_route" : pre_route_slew_ff_rise,
                "SS_rise_slew_pre_route" : pre_route_slew_ss_rise, 
                "FF_fall_slew_pre_route" : pre_route_slew_ff_fall,
                "SS_fall_slew_pre_route" : pre_route_slew_ss_fall,
            }
        else:
            pins_no_appear_in_sdf.append(pin_full_name)
    #print(f"no appear pin in sdf : {pins_no_appear_in_sdf}")
    return pin_info

def generate_graph(def_file, lib_dir, lef_dir, timing_rpt_dir, sdf_file_path, pre_route_sdf_path, graph_save_path):
    pin_info = parse_pins(def_file,lib_dir, timing_rpt_dir,lef_dir, sdf_file_path, pre_route_sdf_path)
    sdf_file_info ,net_connections= parse_sdf_net_section(sdf_file_path)
    cell_connections = parse_sdf_cell_connections(sdf_file_path)

    pin_to_id = {pin: i for i, pin in enumerate(pin_info.keys())}
    id_to_pin = {i: pin for pin, i in pin_to_id.items()}

    # Step 2: 提取节点特征
    feature_list = []
    is_endpoint_list = []
    net_delay_list = []
    for pin in pin_info.keys():
        features = [
            pin_info[pin]["is_PI_PO_or_not"],
            pin_info[pin]["is_fanin_or_fanout"],
            pin_info[pin]["left_distance"],
            pin_info[pin]["right_distance"],
            pin_info[pin]["top_distance"],
            pin_info[pin]["bottom_distance"],
            pin_info[pin]["ff_rise"],
            pin_info[pin]["ff_fall"],
            pin_info[pin]["ss_rise"],
            pin_info[pin]["ss_fall"]
        ]
        feature_list.append(features)
        is_endpoint_list.append(pin_info[pin]["is_endpoint"])
        net_delay = [
            pin_info[pin]["FF_rise_net_delay_to_root_pin"],
            pin_info[pin]["SS_rise_net_delay_to_root_pin"],
            pin_info[pin]["FF_fall_net_delay_to_root_pin"],
            pin_info[pin]["SS_fall_net_delay_to_root_pin"],
        ]
        net_delay_list.append(net_delay)

    node_features = torch.tensor(feature_list, dtype=torch.float32)  # g.ndata['nf']
    is_endpoint = torch.tensor(is_endpoint_list, dtype=torch.float32) # g.ndata['n_is_timing_endpt']
    net_delay_to_root_pin_list = torch.tensor(net_delay_list, dtype=torch.float32)
    # Step 3: 构建 net_out 边
    net_edges = [(pin_to_id[src], pin_to_id[dst]) for src, dst in net_connections]
    net_src_nodes, net_dst_nodes = zip(*net_edges)

    net_out_edge_features = []
    for src_pin, tgt_pin in net_connections:
        # 获取源节点和目标节点的 x_pos 和 y_pos
        src_x_pos = pin_info[src_pin]["left_distance"]
        src_y_pos = pin_info[src_pin]["bottom_distance"]
        tgt_x_pos = pin_info[tgt_pin]["left_distance"]
        tgt_y_pos = pin_info[tgt_pin]["bottom_distance"]

        # 计算 x 和 y 方向上的距离差
        distance_x = tgt_x_pos - src_x_pos
        distance_y = tgt_y_pos - src_y_pos

        # 将距离作为边特征
        net_out_edge_features.append((distance_x, distance_y))

    net_out_edge_features_tensor = torch.tensor(net_out_edge_features, dtype=torch.float32)
    #net_in
    net_in_edges = [(pin_to_id[dst], pin_to_id[src]) for src, dst in net_connections]  # 反向边
    net_in_src_nodes, net_in_dst_nodes = zip(*net_in_edges)
    net_in_edge_features_tensor = -net_out_edge_features_tensor

    # Step 4: 构建 cell_out 边
    cell_edges = []
    for src, dst in cell_connections:
        # 检查 src 和 dst 是否在 pin_info 中
        if src in pin_to_id and dst in pin_to_id:
            # 如果都存在，则添加连接
            cell_edges.append((pin_to_id[src], pin_to_id[dst]))
        else:
            # 如果有缺失的引脚名称，打印警告信息
            if src not in pin_to_id:
                print(f"警告: 引脚 '{src}' 未在 pin_info 中定义")
            if dst not in pin_to_id:
                print(f"警告: 引脚 '{dst}' 未在 pin_info 中定义")

    cell_src_nodes, cell_dst_nodes = zip(*cell_edges)

    # Step 5: 创建 DGL 图
    g = dgl.heterograph({
        ('pin', 'net_out', 'pin'): (net_src_nodes, net_dst_nodes),
        ('pin', 'net_in', 'pin'): (net_in_src_nodes, net_in_dst_nodes),
        ('pin', 'cell_out', 'pin'): (cell_src_nodes, cell_dst_nodes)
    })

    # Step 6: 将节点特征添加到图中
    g.ndata['nf'] = node_features  # 节点特征
    g.ndata['n_is_timing_endpt'] = is_endpoint  # 是否是 endpoint
    g.ndata['n_net_delays'] = net_delay_to_root_pin_list

    g.edges['net_out'].data['ef'] = net_out_edge_features_tensor
    g.edges['net_in'].data['ef'] = net_in_edge_features_tensor

    # Step 7: 打印图信息
    print(g.etypes)
    print("图的节点数:", g.number_of_nodes())
    print("图的边数:", g.number_of_edges())
    print("节点特征维度 (nf):", g.ndata['nf'].shape)
    print("节点特征维度 (n_is_timing_endpt):", g.ndata['n_is_timing_endpt'].shape)
    num_zeros = (g.ndata['n_is_timing_endpt'] > 0.5).sum().item()

    print(f"n_is_timing_endpt 中等于 0 的个数: {num_zeros}")
    print(g.ndata['n_net_delays'].shape)
    print("节点特征 (nf):", g.ndata['nf'])
    print("net_out 边的特征 (ef):", g.edges['net_out'].data['ef'])
    print("net_in 边的特征 (ef):", g.edges['net_in'].data['ef'])
    print(g.edges['net_out'].data['ef'].shape)
    print(g.edges['net_in'].data['ef'].shape)

    #保存图
    dgl.save_graphs(graph_save_path, [g])

def generate_graph_for_dataset(dataset_name, dataset_dir, lib_dir, lef_dir, graph_output_dir, pdk_name):
    log_dir = f"{dataset_dir}/{dataset_name}/{pdk_name}/innovus/log"
    def_file_path = f"{log_dir}/cts.def"
    timing_rpt_dir = log_dir
    sdf_file_path = f"{log_dir}/post_route.sdf"
    pre_route_sdf_path = sdf_file_path = f"{log_dir}/pre_route.sdf"
    graph_output_path = f"{graph_output_dir}/{dataset_name}.graph.bin"

    generate_graph(def_file_path, lib_dir, lef_dir, timing_rpt_dir, sdf_file_path, pre_route_sdf_path, graph_output_path)

def main():
    lib_dir = "/home/hujiahui/case_work/prase_platfrom/asap7/lib"
    lef_dir = "/home/hujiahui/case_work/platform/asap7/lef"
    graph_output_dir = "/home/hujiahui/backend/paper/graph"
    dataset_dir = "/home/hujiahui/test/dataset"
    #dataset_dir = "/home/hujiahui/git/paper/dataset"
    pdk_name = 'asap7'

    # 获取数据集名称列表
    dataset_names = os.listdir(dataset_dir)

    # 计算数据集个数和 CPU 核心数
    num_datasets = len(dataset_names)
    num_cpus = multiprocessing.cpu_count()

    # 设置并行进程数，取数据集个数和 CPU 核心数中的较小值
    num_processes = min(num_datasets, num_cpus)

    # 使用多进程处理数据集
    with multiprocessing.Pool(processes=num_processes) as pool:
        pool.starmap(generate_graph_for_dataset, [(dataset_name, dataset_dir, lib_dir, lef_dir, graph_output_dir, pdk_name) for dataset_name in dataset_names])


def test():
        sdf_file = "/home/hujiahui/test/dataset/aes_cipher/asap7/innovus/log/opensta_pre_route.sdf"
        print(parse_pre_route_sdf(sdf_file))
        

if __name__ == "__main__":
    test()
    #main()


